import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical

class ValueEstimator(nn.Module):
    """价值评估网络，用于估计状态价值"""
    
    def __init__(self, state_dim: int, hidden_dim: int = 128):
        """
        初始化价值评估器
        :param state_dim: 状态维度
        :param hidden_dim: 隐藏层维度
        """
        super().__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, 1)  # 输出状态价值
        
        # 初始化权重
        self._initialize_weights()
    
    def _initialize_weights(self):
        """初始化网络权重"""
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.orthogonal_(m.weight)
                nn.init.zeros_(m.bias)
    
    def forward(self, state: torch.Tensor) -> torch.Tensor:
        """
        前向传播
        :param state: 状态张量 (batch_size, state_dim)
        :return: 状态价值 (batch_size, 1)
        """
        x = F.relu(self.fc1(state))
        x = F.relu(self.fc2(x))
        value = self.fc3(x)
        return value
    
    def estimate(self, state: torch.Tensor) -> float:
        """
        估计单个状态的价值
        :param state: 状态张量 (state_dim,)
        :return: 价值估计值
        """
        with torch.no_grad():
            state = state.unsqueeze(0)  # 添加批次维度
            value = self.forward(state)
            return value.item()
    
    def compute_loss(self, predicted_values: torch.Tensor, target_values: torch.Tensor) -> torch.Tensor:
        """
        计算价值损失（MSE）
        :param predicted_values: 预测价值
        :param target_values: 目标价值（优势值）
        :return: 损失值
        """
        return F.mse_loss(predicted_values, target_values)